package com.yth.utils;

import org.slf4j.Logger;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * Thread utils.
 *
 * @author yutianhong
 * @version 1.0
 * @since 2024/11/27 10:53
 */
public final class ThreadUtils {

    private static final int DEFAULT_THREAD_MULTIPLIER = 2;
    private static final int DEFAULT_RETRIES = 3;
    private static final long DEFAULT_AWAIT_TERMINATION_MILLIS = 100L;

    /**
     * Calculate a suitable number of threads using the default thread multiplier.
     * The result is typically 1.5-2 times the number of CPU cores.
     *
     * @return Suitable thread count
     */
    public static int getSuitableThreadCount() {
        return getSuitableThreadCount(DEFAULT_THREAD_MULTIPLIER);
    }

    /**
     * Calculate a suitable number of threads using the specified thread multiplier.
     *
     * @param threadMultiplier The multiplier for the number of cores
     * @return Suitable thread count
     */
    public static int getSuitableThreadCount(int threadMultiplier) {
        if (threadMultiplier <= 0) {
            throw new IllegalArgumentException("Thread multiplier must be greater than zero");
        }
        int coreCount = Runtime.getRuntime().availableProcessors();
        int targetThreadCount = coreCount * threadMultiplier;

        // Find the next power of two greater than or equal to the target thread count
        return nextPowerOfTwo(targetThreadCount);
    }

    public static void shutdownThreadPool(ExecutorService executor) {
        shutdownThreadPool(executor, null, DEFAULT_RETRIES, DEFAULT_AWAIT_TERMINATION_MILLIS);
    }

    public static void shutdownThreadPool(ExecutorService executor, Logger logger) {
        shutdownThreadPool(executor, logger, DEFAULT_RETRIES, DEFAULT_AWAIT_TERMINATION_MILLIS);
    }

    /**
     * Shutdown thread pool.
     *
     * @param executor                 the thread pool to shut down
     * @param logger                   the logger for logging errors (nullable)
     * @param maxRetries               the maximum number of retries to await termination
     * @param waitTimeBetweenRetriesMs the wait time between retries in milliseconds
     */
    public static void shutdownThreadPool(ExecutorService executor, Logger logger,
                                          int maxRetries, long waitTimeBetweenRetriesMs) {
        if (executor == null) {
            throw new IllegalArgumentException("ExecutorService cannot be null");
        }
        executor.shutdown();
        for (int retry = 0; retry < maxRetries; retry++) {
            try {
                if (executor.awaitTermination(waitTimeBetweenRetriesMs, TimeUnit.MILLISECONDS)) {
                    return;
                }
            } catch (InterruptedException e) {
                executor.shutdownNow();
                Thread.currentThread().interrupt();
                return;
            } catch (Exception ex) {
                if (logger != null) {
                    logger.error("Error occurred while shutting down executor: ", ex);
                }
            }
        }
        executor.shutdownNow();
    }

    /**
     * Adds a shutdown hook to the runtime.
     *
     * @param runnable the runnable to execute upon shutdown
     */
    public static void addShutdownHook(Runnable runnable) {
        Runtime.getRuntime().addShutdownHook(new Thread(runnable));
    }

    /**
     * Calculate the next power of two greater than or equal to the specified number.
     *
     * @param number The number to find the next power of two for
     * @return The next power of two
     */
    private static int nextPowerOfTwo(int number) {
        if (number <= 0) {
            return 1;
        }
        return Integer.highestOneBit(number - 1) << 1;
    }
}
